package com.sicau.rag;

import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;

/**
 * 创建基于向量数据库的文档检索增强顾问工厂
 */
public class RAGAdvisorFactory {
    public static Advisor createRetrievalAugmentationAdvisor(VectorStore vectorStore) {
        // 构造基于向量数据库的文档检索器
        VectorStoreDocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.5)
                .topK(5)
                .build();
        // 构造基于文档检索的增强器
        ContextualQueryAugmenter queryAugmenter = ContextualQueryAugmenter.builder()
                .allowEmptyContext(false)
                .emptyContextPromptTemplate(new PromptTemplate(
                        """
                        请使用以下内容回答用户：
                        对不起，我只能回答心理相关问题，如有其他问题，请发送问题到我们的邮箱xxx@163.com。
                        """
                ))
                .build();
        return RetrievalAugmentationAdvisor.builder()
                .documentRetriever(documentRetriever)
                .queryAugmenter(queryAugmenter)
                .build();
    }
}
